"""
Copyright (c) 2023 Algebra-FUN 
All rights reserved.
"""

import atexit
import io
import re

import cv2 as cv
import numpy as np
import pytesseract
from DrissionPage import ChromiumOptions, ChromiumPage
from DrissionPage._units.listener import DataPacket
from DrissionPage.items import ChromiumElement
from PIL import Image


def simple_binary_preprocess(image: Image.Image):
    gray = cv.cvtColor(np.array(image), cv.COLOR_BGR2GRAY)
    _, binary = cv.threshold(gray, 0, 255, cv.THRESH_BINARY_INV | cv.THRESH_OTSU)
    return Image.fromarray(binary)


def recognize_captcha(captcha_elem: ChromiumElement) -> str:
    img_bytes = captcha_elem.get_screenshot(as_bytes="png")
    image = Image.open(io.BytesIO(img_bytes), formats=["png"])
    image = simple_binary_preprocess(image)
    return pytesseract.image_to_string(image).strip().replace(" ", "")


class Bot:
    def __init__(self, base_url: str = "https://www.insurcloud.com.cn/"):
        self.base_url = base_url
        self.option = ChromiumOptions()
        self.option.set_argument("--disable-features", "CalculateNativeWinOcclusion")
        self.page = None

    def headless(
        self,
        ua: str = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36",
    ):
        self.option.headless()
        self.option.set_argument("--window-size", "1920,1080")
        self.option.set_user_agent(ua)
        return self

    def start(self):
        self.page = ChromiumPage(self.option)
        atexit.register(self.page.quit)
        self.tab = self.page.latest_tab
        self.tab.get(self.base_url)
        self.tab.set.auto_handle_alert()
        self.tab.listen.start("www.insurcloud.com.cn/ipq/insurance.do")
        return self

    def fill_captcha(self, pat: str = "[a-zA-Z0-9]{4}"):
        repat = re.compile(pat)
        captcha_elem = self.tab.ele("css:img#captcha")
        captcha_elem.click(by_js=True)
        answer = recognize_captcha(captcha_elem)
        if repat.fullmatch(answer):
            self.tab.ele("css:input#yzm").input(answer, clear=True, by_js=True)
            return answer
        return None

    def search_one(self, param: dict[str, str], remain_retry: int = 10):
        if remain_retry <= 0:
            return None

        if self.fill_captcha() is None:
            return self.search_one(param, remain_retry=remain_retry - 1)

        for k, v in param.items():
            self.tab.ele(f"css:input#{k}").input(v, clear=True, by_js=True)

        self.tab.listen.resume()
        self.tab.ele("css:#chaxun").click(by_js=True)
        res = self.tab.listen.wait()
        self.tab.listen.pause()

        self.page.close_tabs(self.tab, others=True)

        if isinstance(res, DataPacket):
            body = res.response.body
            if isinstance(body, dict) and body.get("result", "fail") == "succeed":
                return body["dataGrid"]["rows"]

        return self.search_one(param, remain_retry=remain_retry - 1)

    def quit(self):
        if self.page:
            atexit.unregister(self.page.quit)
            self.page.quit()

    def __del__(self):
        self.quit()
